#!/bin/bash

# BBR 自动配置脚本（修复TFO版本）
# 功能：自动开启 BBR 拥塞控制算法并验证

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查 root 权限
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要 root 权限运行，请使用 sudo 执行"
        exit 1
    fi
}

# 检查内核版本
check_kernel_version() {
    local kernel_version=$(uname -r | cut -d. -f1-2)
    local major=$(echo $kernel_version | cut -d. -f1)
    local minor=$(echo $kernel_version | cut -d. -f2)
    
    if [ $major -lt 4 ] || ([ $major -eq 4 ] && [ $minor -lt 9 ]); then
        log_error "内核版本过低 ($(uname -r))，BBR 需要 Linux 4.9+"
        log_info "请升级内核后再运行此脚本"
        exit 1
    else
        log_info "内核版本检查通过: $(uname -r)"
    fi
}

# 检查当前 BBR 状态
check_current_bbr() {
    local current_cc=$(sysctl -n net.ipv4.tcp_congestion_control 2>/dev/null || echo "unknown")
    local current_tfo=$(sysctl -n net.ipv4.tcp_fastopen 2>/dev/null || echo "unknown")
    
    if [ "$current_cc" = "bbr" ] && [ "$current_tfo" = "3" ]; then
        log_info "BBR 和 TFO 已经正确配置"
        return 0
    else
        log_info "当前拥塞控制: $current_cc"
        log_info "当前 TCP Fast Open: $current_tfo (推荐值: 3)"
        return 1
    fi
}

# 获取推荐的队列规则
get_recommended_qdisc() {
    # 检查内核版本，6.0+ 推荐使用 fq_codel
    local kernel_major=$(uname -r | cut -d. -f1)
    if [ "$kernel_major" -ge 6 ]; then
        echo "fq_codel"
    else
        echo "fq"
    fi
}

# 备份原有配置
backup_config() {
    local backup_file="/etc/sysctl.conf.backup.$(date +%Y%m%d_%H%M%S)"
    if cp /etc/sysctl.conf "$backup_file"; then
        log_info "已备份原配置到: $backup_file"
    else
        log_error "备份配置失败"
        exit 1
    fi
}

# 配置 BBR 和 TFO
configure_bbr() {
    log_info "开始配置 BBR 和 TCP Fast Open..."
    
    local recommended_qdisc=$(get_recommended_qdisc)
    log_info "推荐队列规则: $recommended_qdisc (基于内核版本 $(uname -r))"
    
    # 移除可能存在的旧配置
    sed -i '/^# BBR Congestion Control/d' /etc/sysctl.conf
    sed -i '/^net.core.default_qdisc/d' /etc/sysctl.conf
    sed -i '/^net.ipv4.tcp_congestion_control/d' /etc/sysctl.conf
    sed -i '/^net.ipv4.tcp_fastopen/d' /etc/sysctl.conf
    sed -i '/^# Network buffer optimization/d' /etc/sysctl.conf
    sed -i '/^net.core.rmem_max/d' /etc/sysctl.conf
    sed -i '/^net.core.wmem_max/d' /etc/sysctl.conf
    sed -i '/^net.core.rmem_default/d' /etc/sysctl.conf
    sed -i '/^net.core.wmem_default/d' /etc/sysctl.conf
    sed -i '/^net.ipv4.tcp_rmem/d' /etc/sysctl.conf
    sed -i '/^net.ipv4.tcp_wmem/d' /etc/sysctl.conf
    
    # 添加 BBR 和 TFO 配置
    cat >> /etc/sysctl.conf << EOF

# BBR Congestion Control - Auto configured by script
net.core.default_qdisc = $recommended_qdisc
net.ipv4.tcp_congestion_control = bbr
net.ipv4.tcp_fastopen = 3

# Network buffer optimization
net.core.rmem_max = 67108864
net.core.wmem_max = 67108864
net.core.rmem_default = 65536
net.core.wmem_default = 65536
net.ipv4.tcp_rmem = 4096 87380 67108864
net.ipv4.tcp_wmem = 4096 65536 67108864
EOF

    log_info "BBR 和 TCP Fast Open 配置已写入 sysctl.conf"
}

# 应用配置
apply_config() {
    log_info "应用新的系统配置..."
    
    # 静默应用配置，忽略不存在的参数警告
    if sysctl -p > /dev/null 2>&1; then
        log_info "系统配置应用成功"
    else
        log_warn "应用配置时出现一些警告（某些参数可能不存在），正在显示详细输出..."
        sysctl -p
    fi
}

# 验证 BBR 和 TFO
verify_bbr() {
    log_info "验证 BBR 和 TCP Fast Open 配置..."
    
    local current_cc=$(sysctl -n net.ipv4.tcp_congestion_control 2>/dev/null)
    local current_qdisc=$(sysctl -n net.core.default_qdisc 2>/dev/null)
    local current_tfo=$(sysctl -n net.ipv4.tcp_fastopen 2>/dev/null || echo "unknown")
    local bbr_module=$(lsmod | grep bbr | head -1)
    
    echo "==========================================="
    
    # 检查 BBR
    if [ "$current_cc" = "bbr" ]; then
        log_info "✅ BBR 启用成功！"
        log_info "   拥塞控制算法: $current_cc"
    else
        log_error "❌ BBR 启用失败"
        log_info "当前拥塞控制: $current_cc"
        exit 1
    fi
    
    # 检查队列规则
    local valid_qdiscs=("fq" "fq_codel" "fq_pacing" "cake")
    local is_valid_qdisc=0
    for qdisc in "${valid_qdiscs[@]}"; do
        if [ "$current_qdisc" = "$qdisc" ]; then
            is_valid_qdisc=1
            break
        fi
    done
    
    if [ $is_valid_qdisc -eq 1 ]; then
        log_info "   ✅ 队列规则有效: $current_qdisc"
    else
        log_warn "   ⚠️  队列规则可能不是最优: $current_qdisc"
    fi
    
    # 检查 TFO
    if [ "$current_tfo" = "3" ]; then
        log_info "   ✅ TCP Fast Open 已启用 (模式: $current_tfo)"
    else
        log_error "   ❌ TCP Fast Open 配置失败: $current_tfo (应该是 3)"
        log_info "   尝试手动修复..."
        sysctl -w net.ipv4.tcp_fastopen=3
        local fixed_tfo=$(sysctl -n net.ipv4.tcp_fastopen 2>/dev/null || echo "unknown")
        if [ "$fixed_tfo" = "3" ]; then
            log_info "   ✅ 已手动修复 TCP Fast Open"
        else
            log_error "   ❌ 无法修复 TCP Fast Open"
        fi
    fi
    
    if [ -n "$bbr_module" ]; then
        log_info "   ✅ BBR 模块已加载"
    fi
    
    # 显示当前连接状态示例
    log_info "当前 TCP 连接状态示例:"
    ss -tin 2>/dev/null | head -5 || log_warn "无法获取 TCP 连接状态"
    
    echo "==========================================="
    log_info "BBR 配置完成！重启后依然有效"
}

# 显示系统信息
show_system_info() {
    log_info "系统信息:"
    log_info "  - 内核版本: $(uname -r)"
    log_info "  - 操作系统: $(grep PRETTY_NAME /etc/os-release | cut -d= -f2 | tr -d '\"')"
}

# 主函数
main() {
    echo "==========================================="
    log_info "开始 BBR 自动配置"
    echo "==========================================="
    
    show_system_info
    echo ""
    
    check_root
    check_kernel_version
    
    if check_current_bbr; then
        log_info "系统已经正确配置 BBR 和 TFO，无需重复配置"
        verify_bbr
        exit 0
    fi
    
    backup_config
    configure_bbr
    apply_config
    verify_bbr
}

# 执行主函数
main "$@"